import cupy
import gzip
files={
    "x_train":"train-images-idx3-ubyte.gz",
    "t_train":"train-labels-idx1-ubyte.gz",
    "x_test":"t10k-images-idx3-ubyte.gz",
    "t_test":"t10k-labels-idx1-ubyte.gz"
}
def load_label(filename):
    with gzip.open(filename,mode='rb') as f:
        label=cupy.frombuffer(f.read(),count=-1,offset=8,dtype=cupy.uint8)#标签数据偏移8字节
    #将标签转化为独热向量
    hot_vector=cupy.zeros((label.shape[0],10))
    hot_vector[range(label.shape[0]),label]=1
    return hot_vector
def load_feature(filename):
    with gzip.open(filename,mode='rb') as f:
        features=cupy.frombuffer(f.read(),count=-1,offset=16,dtype=cupy.uint8)#图像数据偏移16字节
    features=features.reshape((-1,784))/255
    return features
def relu(x,w,b):
    return cupy.max(cupy.stack((cupy.dot(x,w.T)+b,cupy.zeros((x.shape[0],b.shape[0])))),axis=0)
def I(w,x,b):
    return cupy.where((cupy.dot(x,w)+b)>0,cupy.array(1),cupy.array(0))
def softmax(w1,w2,x,b1,b2,s):#计算属于s类别的概率
    xI = relu(x, w1, b1)
    return cupy.exp(cupy.dot(xI,w2[s,:])+b2[s])/cupy.sum(cupy.exp(cupy.dot(xI,w2.T)+b2),axis=1)
#w1隐藏层，w2输出层
def update(w1,w2,b1,b2,data,label,epoch,n1,n2,batch,test,test_label,eta):
    number=int(len(data)/batch)
    for i in range(epoch):
        for num in range(number):
            #随机选择小批量
            batch_index=cupy.random.randint(number)
            batch_data=data[batch_index*batch:batch_index*batch+batch,:]
            batch_label=label[batch_index*batch:batch_index*batch+batch,:]
            # new_b2=b2.copy()
            # new_w2=w2.copy()
            for j in range(n2):#对输出层进行更新
                tmp=softmax(w1,w2,batch_data,b1,b2,j)-batch_label[:,j]
                w2[j,:]=w2[j,:]-eta*cupy.dot(tmp,relu(batch_data,w1,b1))/batch
                b2[j]=b2[j]-eta*cupy.sum(tmp)/batch
            #使用输出层新的权重更新隐藏层
            for j in range(n1):#对隐藏层进行更新
                xI=relu(batch_data,w1,b1)#计算一个批量的点
                e=cupy.exp(cupy.dot(xI,w2.T) + b2)
                tmp1=cupy.dot(e,w2[:,j])/cupy.sum(e,axis=1)
                tmp2=cupy.sum((cupy.tile(w2[:,j],(xI.shape[0],1))-tmp1.reshape((xI.shape[0],-1)))*batch_label,axis=1)
                w1[j,:]=w1[j,:]+eta*1/batch*cupy.dot(tmp2,I(w1[j,:],batch_data,b1[j]).reshape((-1,1))*batch_data)
                b1[j]=b1[j]+eta*1/batch*cupy.dot(tmp2,I(w1[j,:],batch_data,b1[j]))
            # w2=new_w2
            # b2=new_b2
            #参数更新一次
            #计算此时在测试集上的正确率
            res1 = cupy.empty((batch, n2))
            res2=cupy.empty((len(test),n2))
            for j in range(n2):
                #计算类别
                res1[:,j]=softmax(w1,w2,batch_data,b1,b2,j)
                res2[:,j]=softmax(w1,w2,test,b1,b2,j)
            #每一行最大值为1，其余为0
            index1=cupy.argmax(res1,axis=1)
            index2=cupy.argmax(res2,axis=1)
            res1=cupy.zeros((batch,n2))
            res2=cupy.zeros((test.shape[0],n2))
            res1[range(batch),index1]=1#按照坐标挑选值
            res2[range(len(test)), index2] = 1
            print(f'训练集正确率为:{cupy.sum(cupy.all(res1==batch_label,axis=1))/batch},测试集正确率为:{cupy.sum(cupy.all(res2==test_label,axis=1))/len(test)}')
    return w1,w2,b1,b2
batch=400
n1=256
n2=10
t_train=load_label(files['t_train'])
x_train=load_feature(files['x_train'])
t_test=load_label(files['t_test'])
x_test=load_feature(files['x_test'])
w1=cupy.random.randn(n1,x_train.shape[1])
b1=cupy.random.randn(n1)
w2=cupy.random.randn(n2,n1)
b2=cupy.random.randn(n2)
w1,w2,b1,b2=update(w1,w2,b1,b2,x_train,t_train,100,n1,n2,batch,x_test,t_test,0.01)
#层数越多每次更新所需的数据量越多，需要增大每次迭代样本的数量